#!/usr/bin/env python
# -*- coding: utf-8 -*-

import json
import os
import random
import datetime
import multiprocessing
from math import sin, cos, radians

def generate_waypoints(starting_point, num_points=150, flight_type="circular"):
    """生成航点列表
    
    Args:
        starting_point: 起始点，格式为 (经度, 纬度, 高度)
        num_points: 生成的航点数量，默认150个
        flight_type: 飞行类型，可选 "circular", "grid", "zigzag", "random"
    
    Returns:
        包含航点的列表
    """
    points = []
    start_lon, start_lat, start_alt = starting_point
    timestamp_base = datetime.datetime(2023, 6, 1, 8, 0, 0)  # 基准时间
    
    # 首先添加起飞阶段（垂直上升）
    takeoff_height = random.randint(100, 200)  # 随机选择飞行高度
    takeoff_steps = random.randint(5, 8)  # 垂直上升的点数
    
    for i in range(takeoff_steps):
        altitude = start_alt + (takeoff_height - start_alt) * (i / takeoff_steps)
        timestamp = timestamp_base + datetime.timedelta(seconds=i*10)
        points.append({
            "longitude": start_lon,
            "latitude": start_lat,
            "altitude": altitude,
            "timestamp": timestamp.strftime("%Y-%m-%dT%H:%M:%SZ")
        })
    
    # 生成巡航阶段的航点
    remaining_points = num_points - takeoff_steps - 5  # 减去起飞和降落的点数
    
    # 根据不同的飞行模式生成航点
    if flight_type == "circular":
        # 圆形飞行模式
        radius = random.uniform(0.005, 0.02)  # 随机半径（约0.5-2公里）
        angle_step = 360 / (remaining_points / 2)  # 角度步长
        
        for i in range(remaining_points):
            angle = radians(i * angle_step)
            # 第一圈
            if i < remaining_points / 2:
                lon = start_lon + radius * cos(angle)
                lat = start_lat + radius * sin(angle)
            # 第二圈（可以有高度变化）
            else:
                lon = start_lon + radius * 0.8 * cos(angle)
                lat = start_lat + radius * 0.8 * sin(angle)
                
            timestamp = timestamp_base + datetime.timedelta(seconds=(i+takeoff_steps)*10)
            points.append({
                "longitude": lon,
                "latitude": lat,
                "altitude": takeoff_height,
                "timestamp": timestamp.strftime("%Y-%m-%dT%H:%M:%SZ")
            })
            
    elif flight_type == "grid":
        # 网格飞行模式
        grid_size = int(remaining_points**0.5)  # 网格大小
        step_size = 0.001 * random.uniform(1, 3)  # 网格步长
        
        for i in range(grid_size):
            # 沿x轴移动
            for j in range(grid_size):
                if i % 2 == 0:
                    # 正向
                    lon = start_lon + j * step_size
                else:
                    # 反向
                    lon = start_lon + (grid_size - 1 - j) * step_size
                    
                lat = start_lat + i * step_size
                timestamp = timestamp_base + datetime.timedelta(seconds=(i*grid_size+j+takeoff_steps)*10)
                
                if i*grid_size+j < remaining_points:
                    points.append({
                        "longitude": lon,
                        "latitude": lat,
                        "altitude": takeoff_height,
                        "timestamp": timestamp.strftime("%Y-%m-%dT%H:%M:%SZ")
                    })
    
    elif flight_type == "zigzag":
        # 之字形飞行模式
        amplitude = random.uniform(0.005, 0.015)  # 振幅
        step_lon = 0.001 * random.uniform(0.5, 1.5)  # 经度步长
        
        for i in range(remaining_points):
            progress = i / remaining_points
            lon = start_lon + progress * step_lon * remaining_points
            lat = start_lat + amplitude * sin(progress * 12)  # 控制之字形的频率
            
            timestamp = timestamp_base + datetime.timedelta(seconds=(i+takeoff_steps)*10)
            
            # 随机高度变化
            alt_variation = random.randint(-20, 20) if i > remaining_points/3 else 0
            
            points.append({
                "longitude": lon,
                "latitude": lat,
                "altitude": takeoff_height + alt_variation,
                "timestamp": timestamp.strftime("%Y-%m-%dT%H:%M:%SZ")
            })
    
    else:  # random
        # 随机飞行模式
        for i in range(remaining_points):
            # 随机偏移，但保持连续性
            if i == 0:
                lon = start_lon + random.uniform(-0.002, 0.002)
                lat = start_lat + random.uniform(-0.002, 0.002)
            else:
                prev_lon = points[takeoff_steps + i - 1]["longitude"]
                prev_lat = points[takeoff_steps + i - 1]["latitude"]
                lon = prev_lon + random.uniform(-0.001, 0.001)
                lat = prev_lat + random.uniform(-0.001, 0.001)
            
            timestamp = timestamp_base + datetime.timedelta(seconds=(i+takeoff_steps)*10)
            
            # 随机高度变化
            if i > remaining_points / 2:
                alt_variation = random.randint(-40, 0)
            else:
                alt_variation = random.randint(-10, 10)
                
            points.append({
                "longitude": lon,
                "latitude": lat,
                "altitude": takeoff_height + alt_variation,
                "timestamp": timestamp.strftime("%Y-%m-%dT%H:%M:%SZ")
            })
    
    # 添加降落阶段
    last_lon = points[-1]["longitude"]
    last_lat = points[-1]["latitude"]
    last_alt = points[-1]["altitude"]
    
    landing_steps = 5
    for i in range(landing_steps):
        progress = (i + 1) / landing_steps
        altitude = last_alt * (1 - progress) + start_alt * progress
        timestamp = timestamp_base + datetime.timedelta(seconds=(remaining_points+takeoff_steps+i)*10)
        
        # 逐渐返回起始位置
        lon = last_lon * (1 - progress) + start_lon * progress
        lat = last_lat * (1 - progress) + start_lat * progress
        
        points.append({
            "longitude": lon,
            "latitude": lat,
            "altitude": altitude,
            "timestamp": timestamp.strftime("%Y-%m-%dT%H:%M:%SZ")
        })
    
    return points

def create_uav_route_file(uav_id, output_dir, starting_point=None):
    """创建单个无人机航线文件
    
    Args:
        uav_id: 无人机ID号，例如34
        output_dir: 输出目录
        starting_point: 起始点，如果为None则随机生成
    """
    if starting_point is None:
        # 随机生成起始点，北京附近
        longitude = random.uniform(116.2, 116.6)
        latitude = random.uniform(39.8, 40.2)
        altitude = random.randint(40, 60)
        starting_point = (longitude, latitude, altitude)
    
    # 随机选择飞行模式
    flight_type = random.choice(["circular", "grid", "zigzag", "random"])
    
    # 生成航点
    num_points = random.randint(150, 200)  # 至少150个点
    points = generate_waypoints(starting_point, num_points, flight_type)
    
    # 创建无人机航线数据，使用4位数字格式以适应1000以上的ID
    uav_id_str = str(uav_id).zfill(4)
    route_data = {
        "id": f"dji-uav-{uav_id_str}",
        "name": f"大疆无人机{uav_id_str}",
        "status": "online",
        "points": points
    }
    
    # 输出到文件
    output_path = os.path.join(output_dir, f"dji-uav-{uav_id_str}.json")
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(route_data, f, ensure_ascii=False, indent=2)
    
    return f"已创建航线文件: {output_path} ({len(points)}个点, 飞行模式: {flight_type})"

def process_uav_batch(batch_info):
    """处理一批无人机，用于多进程执行
    
    Args:
        batch_info: 包含批次信息的元组 (start_id, end_id, output_dir, regions, fixed_points)
    
    Returns:
        处理结果消息
    """
    start_id, end_id, output_dir, regions, fixed_points = batch_info
    results = []
    
    for uav_id in range(start_id, end_id + 1):
        # 为一些无人机使用固定起始点，其余使用随机起始点但基于区域分布
        if uav_id < 34 + len(fixed_points) and uav_id - 34 < len(fixed_points):
            starting_point = fixed_points[uav_id - 34]
        else:
            # 基于区域选择随机起点
            region_idx = (uav_id - 34) % len(regions)
            region_lon, region_lat = regions[region_idx]
            
            # 在区域内随机偏移
            longitude = region_lon + random.uniform(-0.05, 0.05)
            latitude = region_lat + random.uniform(-0.05, 0.05)
            altitude = random.randint(40, 60)
            starting_point = (longitude, latitude, altitude)
            
        result = create_uav_route_file(uav_id, output_dir, starting_point)
        results.append(result)
    
    return f"批次 {start_id}-{end_id} 完成，共处理 {len(results)} 架无人机"
    
def main():
    # 输出目录
    output_dir = "public/json"
    
    # 确保目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 生成起始ID
    start_id = 34  # 从34开始，因为前面已经有33个文件
    end_id = 1033  # 修改为1000架无人机 (1000 + 33 = 1033)
    
    # 为了有效地处理大量无人机，增加更多的固定起始点区域
    fixed_starting_points = [
        # 北京市区及周边区域
        (116.3, 39.9, 50),    # 北京市中心附近
        (116.25, 39.85, 50),  # 西南方向
        (116.45, 40.05, 50),  # 东北方向
        (116.2, 40.1, 50),    # 西北方向
        (116.5, 39.8, 50),    # 东南方向
        # 扩展更多区域
        (116.4, 39.95, 50),   # 市中心偏东
        (116.35, 40.0, 50),   # 市中心偏北
        (116.28, 39.95, 50),  # 市中心偏西
        (116.35, 39.85, 50),  # 市中心偏南
        # 更远区域
        (116.6, 40.1, 50),    # 东北远区
        (116.1, 40.2, 50),    # 西北远区
        (116.7, 39.7, 50),    # 东南远区
        (116.0, 39.75, 50),   # 西南远区
        # 再扩展更多区域以增加分布
        (116.55, 39.95, 50),  # 东部
        (116.15, 39.95, 50),  # 西部
        (116.35, 40.15, 50),  # 北部
        (116.35, 39.75, 50),  # 南部
    ]
    
    # 计算需要处理的无人机总数
    total_uavs = end_id - start_id + 1
    
    print(f"开始生成 {total_uavs} 架无人机的航线数据...")
    
    # 生成区域分布参数，使无人机更均匀分布
    regions = []
    for i in range(0, 4):
        for j in range(0, 4):
            lon_base = 116.1 + i * 0.15
            lat_base = 39.7 + j * 0.15
            regions.append((lon_base, lat_base))
    
    # 使用多进程加速处理
    num_cores = multiprocessing.cpu_count()
    print(f"检测到 {num_cores} 个CPU核心，将使用多进程处理")
    
    # 将任务分成多个批次
    batch_size = max(10, total_uavs // (num_cores * 2))  # 每个批次处理的无人机数量
    batches = []
    
    for batch_start in range(start_id, end_id + 1, batch_size):
        batch_end = min(batch_start + batch_size - 1, end_id)
        batches.append((batch_start, batch_end, output_dir, regions, fixed_starting_points))
    
    print(f"任务已分为 {len(batches)} 个批次进行处理")
    
    # 使用进程池处理所有批次
    with multiprocessing.Pool(processes=num_cores) as pool:
        results = []
        for i, result in enumerate(pool.imap_unordered(process_uav_batch, batches)):
            results.append(result)
            progress = (i + 1) / len(batches) * 100
            print(f"进度: {progress:.1f}% - {result}")
    
    print(f"完成! 已生成 {total_uavs} 架无人机的航线数据。")
    
if __name__ == "__main__":
    main() 